/*
This program detects image features using SIFT keypoints. For more info,
refer to:

Lowe, D. Distinctive image features from scale-invariant keypoints.
International Journal of Computer Vision, 60, 2 (2004), pp.91--110.

Copyright (C) 2006  Rob Hess <hess@eecs.oregonstate.edu>

Note: The SIFT algorithm is patented in the United States and cannot be
used in commercial products without a license from the University of
British Columbia.  For more information, refer to the file LICENSE.ubc
that accompanied this distribution.

Version: 1.1.1-20070913
*/

#include "sift.h"
#include "imgfeatures.h"
#include "utils.h"

#include <highgui.h>
#include <ml.h>

#include "windows.h"
#include <stdio.h>
#include "fstream"
using namespace std;

#define KMEAN_CLUSTER_NUM 20
int kmeanIterationCount = 10; //10;
int kmeanEpsilon = 1; //1.0;

#define FEAT_LEN 128


struct KMeanCenters
{
	int* arrCount;
	float** arrCenters;
	KMeanCenters()
	{
		arrCount = new int[KMEAN_CLUSTER_NUM];
		memset(arrCount,0,KMEAN_CLUSTER_NUM*sizeof(int));
		arrCenters = new float*[KMEAN_CLUSTER_NUM];
		for (int i=0; i<KMEAN_CLUSTER_NUM; i++)
		{
			arrCenters[i] = new float[FEAT_LEN];
			memset(arrCenters[i],0,sizeof(float)*FEAT_LEN);
		}
	}
};

//Prototype
KMeanCenters Kmeans(char* strFolder, int kmeanClusterNum, char* strOutFile);
void trainSVMBagOfWord(char* strPosFold, char* strNegFold, KMeanCenters kCenters, char* strModelFile);
//bool mainPredict(IplImage* img, KMeanCenters kCenters);
bool mainPredict(IplImage* img, KMeanCenters kCenters, char* svmModelFile);
KMeanCenters readClusterFromFile(char* strFileName);
bool SIFT_SVM_Predict(IplImage* img, char* kmeansModel, char* svmModel);
/******************************** Globals ************************************/

char* img_file_name = "..\\trihoang_scale.jpg";//"..\\beaver.png";
char* out_file_name  = "..\\trihoang_scale.sift";;
char* img_referee = "..\\referee.png";
char* img_nonreferee = "..\\nonreferee.png";
char* svm_ModelFile = "SVMModel.txt";
char* img_checkSVM1_1 = "..\\checkSVM1_1.png";
char* img_checkSVM1_2 = "..\\checkSVM1_2.png";
char* img_checkSVM1_3 = "..\\checkSVM1_3.png";
char* img_checkSVM2_1 = "..\\checkSVM2_1.png";
char* img_checkSVM2_2 = "..\\checkSVM2_2.png";
char* img_checkSVM2_3 = "..\\checkSVM2_3.png";
char* img_checkSVM2_4 = "..\\checkSVM2_4.png";
char* svm_PosTrainFold = "E:\\Temporary\\M\\M\\VT_DI1\\Training\\Referee";
char* svm_NegTrainFold = "E:\\Temporary\\M\\M\\VT_DI1\\Training\\Non-Referee";
char* kmean_TrainFold = "E:\\Temporary\\M\\M\\VT_DI1\\Training\\KMEANS";
char* kmean_OutFile = "kmeans.txt";
char* out_img_name = NULL;
int display = 1;
int intvls = SIFT_INTVLS;
double sigma = SIFT_SIGMA;
double contr_thr = SIFT_CONTR_THR;
int curv_thr = SIFT_CURV_THR;
int img_dbl = SIFT_IMG_DBL;
int descr_width = SIFT_DESCR_WIDTH;
int descr_hist_bins = SIFT_DESCR_HIST_BINS;



/***********************************************************************/

int siftFeatureExtract(IplImage *img, feature **feat)
{
	int n = 0;

	if( ! img )
	{
		return 0;
	}
	n = _sift_features( img, (feature**)feat, intvls, sigma, contr_thr, curv_thr,
		img_dbl, descr_width, descr_hist_bins );
	return n;
}


void trainSVMFromMatrix(CvMat *matTrainData, CvMat *matTrainClasses, char* strSVMModelFile)
{

	CvTermCriteria criter=cvTermCriteria(CV_TERMCRIT_EPS,0, 0.0001);
	CvSVMParams param = CvSVMParams(CvSVM::C_SVC, CvSVM::LINEAR, 0, 0, 0, 2,
		0, 0, 0, criter);

	//another param
	//CvSVMParams param;	
	//param.term_crit.epsilon = 0.000001;
	//param.term_crit.max_iter = 100;
	//param.term_crit.type = CV_TERMCRIT_ITER|CV_TERMCRIT_EPS;
	//param.svm_type = CvSVM::NU_SVC;
	//param.kernel_type = CvSVM::RBF;
	//param.gamma = 64;
	//param.C = 8;
	//param.nu = 0.5;

	CvSVM svm(matTrainData, matTrainClasses, 0, 0,
		param);

	svm.save(strSVMModelFile);

	//cvReleaseMat(&trainClasses);
	//cvReleaseMat(&trainData);
}

//void main()
//{
//	//KMeanCenters kCenters;
//	//kCenters = Kmeans(kmean_TrainFold, KMEAN_CLUSTER_NUM, kmean_OutFile);
//	//kCenters = readClusterFromFile(kmean_OutFile);
//	//trainSVMBagOfWord(svm_PosTrainFold,svm_NegTrainFold,kCenters, svm_ModelFile);
//
//	IplImage* img;
//	img = cvLoadImage(img_checkSVM1_1,1);
//	bool result = SIFT_SVM_Predict(img);
//	if (result)
//		printf("la trong tai (T)\n");
//	else
//		printf("khong phai trong tai (T)\n");
//
//	img = cvLoadImage(img_checkSVM1_2,1);
//	result = SIFT_SVM_Predict(img);
//	if (result)
//		printf("la trong tai (T)\n");
//	else
//		printf("khong phai trong tai (T)\n");
//
//	img = cvLoadImage(img_checkSVM1_3,1);
//	result = SIFT_SVM_Predict(img);
//	if (result)
//		printf("la trong tai (T)\n");
//	else
//		printf("khong phai trong tai (T)\n");
//
//	img = cvLoadImage(img_checkSVM2_1,1);
//	result = SIFT_SVM_Predict(img);
//	if (result)
//		printf("la trong tai (K)\n");
//	else
//		printf("khong phai trong tai (K)\n");
//
//	img = cvLoadImage(img_checkSVM2_2,1);
//	result = SIFT_SVM_Predict(img);
//	if (result)
//		printf("la trong tai (K)\n");
//	else
//		printf("khong phai trong tai (K)\n");
//
//	img = cvLoadImage(img_checkSVM2_3,1);
//	result = SIFT_SVM_Predict(img);
//	if (result)
//		printf("la trong tai (K)\n");
//	else
//		printf("khong phai trong tai (K)\n");
//
//	img = cvLoadImage(img_checkSVM2_4,1);
//	result = SIFT_SVM_Predict(img);
//	if (result)
//		printf("la trong tai (K)\n");
//	else
//		printf("khong phai trong tai (K)\n");
//
//}



//Bag of word version's functions
//
//

KMeanCenters Kmeans(char* strFolder, int kmeanClusterNum, char* strOutFile)
{
	//mang cua cac mang feat
	feature* arrFeat[500];
	//so phan tu trong moi mang
	int arrFeaNum[500];
	//so luong mang cac feat
	int iFeaCount = 0;
	//tong so luong feat
	int iFeaSum = 0;


	char* dirName = strFolder;
	char* dirScan = new char[100];
	strcpy(dirScan, strFolder);
	strcat(dirScan, "\\*.jpg");
	char* pname = 0; /* File name buffer */
	WIN32_FIND_DATAA info;
	HANDLE h = FindFirstFileA(dirScan,&info);
	if (h == INVALID_HANDLE_VALUE)
	{
		/* print message?.. */
		KMeanCenters tmp;
		return tmp;
	}
	pname = (char*)malloc(MAX_PATH); /* (char*) for C++ only */
	do
	{
		if ((info.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) != 0)
			continue; /* skip directories */
		strcpy(pname,dirName);
		strcat(pname,"\\");
		strcat(pname,info.cFileName);
		/* Now we have the next file name... */
		int iFeaNumTmp;
		IplImage* img;
		img = cvLoadImage( pname, 1 );
		iFeaNumTmp = siftFeatureExtract(img,&arrFeat[iFeaCount]);
		if (iFeaNumTmp>0)
		{
			arrFeaNum[iFeaCount] = iFeaNumTmp;
			iFeaCount++;
			iFeaSum += iFeaNumTmp;
		}
		printf("%d\n",iFeaCount);

		/* Open, process then close the file ...*/
	} while (FindNextFileA(h,&info));
	FindClose(h); /* Close dir scan */
	//free(pname);


	//chuan bi du lieu cho kmeans
	CvMat* matData = cvCreateMat(iFeaSum,
		arrFeat[0][0].d, CV_32FC1);
	int count = 0;
	for (int i=0; i<iFeaCount; i++)
	{
		for (int j=0; j<arrFeaNum[i]; j++)
		{
			memcpy(matData->data.fl+count*arrFeat[0][0].d,arrFeat[i][j].descr,sizeof(float)*arrFeat[0][0].d);
			count++;
		}
	}

	CvMat* matClusters = cvCreateMat(iFeaSum, 1, CV_32SC1);
	cvKMeans2(matData, KMEAN_CLUSTER_NUM, matClusters, 
		cvTermCriteria(CV_TERMCRIT_EPS+CV_TERMCRIT_ITER, 
		kmeanIterationCount, kmeanEpsilon));

	printf("kmeans completed\n");
	KMeanCenters kCenters;
	for (int i=0; i<iFeaSum; i++)
	{
		int j = (int)cvGetReal1D(matClusters,i);
		if (j>=KMEAN_CLUSTER_NUM)
			printf("over\n");
		kCenters.arrCount[j]++;
		for (int k = 0; k<FEAT_LEN; k++)
		{
			kCenters.arrCenters[j][k]+=cvmGet(matData,i,k);
			if (j==0)
			{
				float ftmp = cvmGet(matData,i,k);
				printf("%f A ",ftmp);
			}
		}
		if (j==0)
			printf("END\n");
	}	
	for (int i=0; i<KMEAN_CLUSTER_NUM; i++)
	{
		if (kCenters.arrCount[i]>0)
		{
			for (int j=0; j<FEAT_LEN; j++)
				kCenters.arrCenters[i][j] = kCenters.arrCenters[i][j] / kCenters.arrCount[i];
		}
	}	
	ofstream outfile(strOutFile);
	outfile<<KMEAN_CLUSTER_NUM;
	outfile<<"\n";
	for (int i=0; i<KMEAN_CLUSTER_NUM; i++)
	{
		outfile<<kCenters.arrCount[i]<<" ";
		for (int j=0 ; j<FEAT_LEN; j++)
		{
			outfile<<kCenters.arrCenters[i][j];
			outfile<<" ";
		}
		outfile<<"\n";
	}
	outfile.close();
	return kCenters;
}

float calcMahattanDistance(float* v1, float* v2, int len)
{
	float d = 0;
	for (int i=0; i<len; i++)
		d += abs(v1[i]-v2[i]);
	return d;
}

int findNearestVector(float** arrV, int lenArr, float* v, int lenv, int* arrCount)
{
	float dmin = 1000000;
	float dtmp;
	int imin = 0;
	for (int i=0; i<lenArr; i++)
		if (arrCount[i]>0)
		{
			dtmp = calcMahattanDistance(arrV[i],v,lenv);
			if (dtmp<dmin)
			{
				imin = i;
				dmin = dtmp;
			}
		}
		return imin;
}

void trainSVMBagOfWord(char* strPosFold, char* strNegFold, KMeanCenters kCenters, char* strModelFile)
{
	float arrPosBOWFeat[500][KMEAN_CLUSTER_NUM];
	int iPosFeatNum=0;
	float arrNegBOWFeat[500][KMEAN_CLUSTER_NUM];
	int iNegFeatNum=0;

	//Process Positive Folder
	char* dirName = strPosFold;
	char* dirScan = new char[100];
	strcpy(dirScan, strPosFold);
	strcat(dirScan, "\\*.png");
	char* pname = 0; /* File name buffer */
	WIN32_FIND_DATAA info;
	HANDLE h = FindFirstFileA(dirScan,&info);
	if (h == INVALID_HANDLE_VALUE)
	{
		return;
	}
	pname = (char*)malloc(MAX_PATH); /* (char*) for C++ only */
	do
	{
		if ((info.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) != 0)
			continue; /* skip directories */
		strcpy(pname,dirName);
		strcat(pname,"\\");
		strcat(pname,info.cFileName);
		/* Now we have the next file name... */
		int iFeaCount;
		IplImage* img;
		feature* feat;
		img = cvLoadImage( pname, 1 );
		iFeaCount = siftFeatureExtract(img,&feat);
		memset(arrPosBOWFeat[iPosFeatNum],0,sizeof(float)*KMEAN_CLUSTER_NUM);
		for (int i=0; i<iFeaCount; i++)
		{
			int bin = findNearestVector(kCenters.arrCenters,KMEAN_CLUSTER_NUM, feat[i].descr, FEAT_LEN,kCenters.arrCount);
			arrPosBOWFeat[iPosFeatNum][bin]++;
		}
		//for (int i=0; i<KMEAN_CLUSTER_NUM; i++)
		//	arrPosBOWFeat[iPosFeatNum][i] /= iFeaCount;
		iPosFeatNum++;
		printf("BOW %d\n",iPosFeatNum);
		/* Open, process then close the file ...*/
	} while (FindNextFileA(h,&info));

	FindClose(h); /* Close dir scan */


	//Process Positive Folder
	dirName = strNegFold;
	dirScan = new char[100];
	strcpy(dirScan, strNegFold);
	strcat(dirScan, "\\*.png");
	pname = 0; /* File name buffer */
	//WIN32_FIND_DATA info;
	h = FindFirstFileA(dirScan,&info);
	if (h == INVALID_HANDLE_VALUE)
	{
		/* print message?.. */
		return;
	}
	pname = (char*)malloc(MAX_PATH); /* (char*) for C++ only */
	do
	{
		if ((info.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) != 0)
			continue; /* skip directories */
		strcpy(pname,dirName);
		strcat(pname,"\\");
		strcat(pname,info.cFileName);
		/* Now we have the next file name... */
		int iFeaCount;
		IplImage* img;
		feature* feat;
		img = cvLoadImage( pname, 1 );
		iFeaCount = siftFeatureExtract(img,&feat);
		memset(arrNegBOWFeat[iNegFeatNum],0,sizeof(float)*KMEAN_CLUSTER_NUM);
		for (int i=0; i<iFeaCount; i++)
		{
			int bin = findNearestVector(kCenters.arrCenters,KMEAN_CLUSTER_NUM, feat[i].descr, FEAT_LEN,kCenters.arrCount);
			arrNegBOWFeat[iNegFeatNum][bin]++;
		}
		//for (int i=0; i<KMEAN_CLUSTER_NUM; i++)
		//	arrNegBOWFeat[iNegFeatNum][i] /= iFeaCount;
		iNegFeatNum++;
		printf("BOW %d\n",iNegFeatNum);
		/* Open, process then close the file ...*/
	} while (FindNextFileA(h,&info));

	FindClose(h); /* Close dir scan */
	free(pname);

	CvMat* matTrainData = cvCreateMat(iPosFeatNum + iNegFeatNum,
		KMEAN_CLUSTER_NUM, CV_32FC1);
	int count = 0;
	for (int i=0; i<iPosFeatNum; i++)
	{
		memcpy(matTrainData->data.fl+count*KMEAN_CLUSTER_NUM,arrPosBOWFeat[i],sizeof(float)*KMEAN_CLUSTER_NUM);
		count++;
	}

	for (int i=0; i<iNegFeatNum; i++)
	{
		memcpy(matTrainData->data.fl+count*KMEAN_CLUSTER_NUM,arrNegBOWFeat[i],sizeof(float)*KMEAN_CLUSTER_NUM);
		count++;
	}
	for (int i =0; i<matTrainData->rows; i++)
	{
		CvMat row;
		cvGetRow(matTrainData,&row,i);
		cvNormalize(&row, &row,1,0,CV_L1);
	}
	for (int i=0; i<matTrainData->rows; i++)
	{
		for (int j=0; j<matTrainData->cols; j++)
			printf("%f ",cvmGet(matTrainData,i,j));
		printf("\n");
	}

	//Create svm training classes
	CvMat* matTrainClasses = cvCreateMat(iPosFeatNum+iNegFeatNum,1, CV_32FC1 );
	CvMat matTrainClasses1, matTrainClasses2;
	cvGetRows(matTrainClasses, &matTrainClasses1, 0, iPosFeatNum);
	cvSet(&matTrainClasses1, cvScalar(1));

	cvGetRows(matTrainClasses, &matTrainClasses2, iPosFeatNum, iPosFeatNum+iNegFeatNum);
	cvSet(&matTrainClasses2, cvScalar(2));

	//printf("train data were prepared\n");
	//train
	trainSVMFromMatrix(matTrainData, matTrainClasses, strModelFile);
	printf("train completed\n");

}


float predictSVMBagOfWord(float* vBOW, char* strSVMModelFile)
{
	CvSVM svm;
	svm.load(strSVMModelFile);
	CvMat* data = cvCreateMat(1,
		KMEAN_CLUSTER_NUM, CV_32FC1);
	memcpy(data->data.fl,vBOW,sizeof(float)*KMEAN_CLUSTER_NUM);
	cvNormalize(data, data,1,0,CV_L1);
	return svm.predict(data);
}

bool mainPredict(IplImage* img, KMeanCenters kCenters, char* svmModelFile)
{
	feature* feat;
	float vBOWFeat[KMEAN_CLUSTER_NUM];
	memset(vBOWFeat,0,sizeof(float)*KMEAN_CLUSTER_NUM);
	int iFeaCount = siftFeatureExtract(img,&feat);
	for (int i=0; i<iFeaCount; i++)
	{
		int bin = findNearestVector(kCenters.arrCenters,KMEAN_CLUSTER_NUM, feat[i].descr, FEAT_LEN, kCenters.arrCount);
		vBOWFeat[bin]++;
	}

	//for (int i=0; i<KMEAN_CLUSTER_NUM; i++)
	//	vBOWFeat[i] /= iFeaCount;
	if (predictSVMBagOfWord(vBOWFeat, svmModelFile/*svm_ModelFile*/) == 1)
		return true;
	else
		return false;
}

KMeanCenters readClusterFromFile(char* strFileName)
{
	KMeanCenters kCenters;
	int n;
	ifstream infile(strFileName);
	infile>>n;
	for (int i=0; i<KMEAN_CLUSTER_NUM; i++)
	{
		infile>>kCenters.arrCount[i];
		for (int j=0; j<FEAT_LEN; j++)
			infile>>kCenters.arrCenters[i][j];
	}
	infile.close();
	return kCenters;
}

//bool SIFT_SVM_Predict(IplImage* img)
bool SIFT_SVM_Predict(IplImage* img, char* kmeansModel, char* svmModel)
{
	KMeanCenters kCenters = readClusterFromFile(kmeansModel/*kmean_OutFile*/);
	return mainPredict(img, kCenters, svmModel);
}